import base64

from DrissionPage._pages.chromium_page import ChromiumPage


class SetRequestResponse:
    def __init__(self, page: ChromiumPage, urls: list) -> None:
        """

        :param page:
        :param urls: 字典数组[{
                            url:地址
                            response:替换成的内容
                            }]
        """
        self._page = page
        self._driver = page.driver
        self.urls = urls
        self.start()

    def start(self):
        self._driver.run("Fetch.enable")
        self._set_callback()

    def _set_callback(self):
        self._driver.set_callback("Fetch.requestPaused", self.response_change)

    def response_change(self, **kwargs):
        request_url = kwargs.get("request").get("url")
        headers = kwargs.get("request").get("headers")
        for url in self.urls:
            if url['url'] in request_url:
                self._driver.run(
                    "Fetch.fulfillRequest",
                    requestId=kwargs.get("requestId"),
                    body=format_Body(url['response'].encode()),
                    responseHeaders=format_Headers(headers),
                    responseCode=200
                )
            else:
                self._driver.run(
                    "Fetch.continueRequest",
                    requestId=kwargs.get("requestId")
                )


def format_Headers(headers):
    _headers = []
    for header in headers:
        _headers.append({"name": header, "value": headers.get(header)})
    return _headers


def format_Body(body):
    base64_body = base64.b64encode(body).decode("utf-8")
    return base64_body


class SetRequestParam:
    def __init__(self, page: ChromiumPage, postData: list) -> None:
        """

        :param page:
        :param infos: 字典数组[{
                            url:地址
                            param:参数
                            }]
        """
        self._page = page
        self._driver = page.driver
        self.postData = postData
        self.start()

    def start(self):
        self._driver.run("Fetch.enable")
        self._set_callback()

    def _set_callback(self):
        self._driver.set_callback("Fetch.requestPaused", self.response_change)

    def response_change(self, **kwargs):
        request_url = kwargs.get("request").get("url")
        for data in self.postData:
            if data['url'] in request_url:
                self._driver.run(
                    "Fetch.continueRequest",
                    requestId=kwargs.get("requestId"),
                    postData=format_Body(data['param'].encode())
                )